import json,re
import websockets,traceback
from src import mj_util
from http.server import HTTPServer, BaseHTTPRequestHandler
from socketserver import ThreadingMixIn
from src import server_handler
from src import config
from src.log import * 
from src.plugin import mj_prompt_optimizer
from concurrent import futures
from src.server_handler import gExecutor


def generateDiscordMsgHash(url):
    parts = url.split('_')
    return parts[-1].split('.')[0]


# 定义HTTP请求处理程序
class MyRequestHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        response_body = 'Internal Error'.encode('utf-8')
        self.send_header('Content-Type', 'application/json')
        self.send_header('Content-Length', len(response_body))
        self.end_headers()
        self.wfile.write(response_body)
    
    def do_POST(self):
        Logger.info(f'do_POST path request:{self.path}')
        # 使用正则表达式匹配路径
        command = re.compile(r'^/midjourney-command$')
        discord_callback = re.compile(r'^/discord-api$')

        # 根据路径执行不同的逻辑
        #处理
        if command.match(self.path):
            content_length = int(self.headers['Content-Length'])
            try:
                body = self.rfile.read(content_length)
                data = json.loads(body)
                Logger.info(f'/midjourney-command$:request-data:{data}')
                type = data.get('type', None)
                msgId = data.get('msgId', None)
                parentMsgId = data.get('parentMsgId', None)
                response_data=dict()
                msg = ''
                if type == 'imagine':
                    prompt = data.get('prompt', None)
                    if prompt is  None or prompt.lstrip().rstrip() == '' or msgId == None:
                        code = -1
                        msg = f'generate:参数无效{prompt} msgId:{msgId}'
                        response_data=dict(code=code, msg=msg)
                    else:
                        gpt_optimization = data.get('gpt_optimization', config.PROMPT_OPTIMZE_FLAG)
                        conversation_id = data.get('conversation_id', '')
                        
                        if gpt_optimization == 1:
                            def next_task(prompt, conversation_id, msgId,parentMsgId):
                                prompt_optmized=mj_prompt_optimizer.optmize_prompt(prompt)
                                id, response_data = server_handler.SendDiscordMessageMJGenerate(conversation_id, msgId, parentMsgId, prompt_optmized)
                                if  response_data['code'] == 0:
                                    pass
                                else:
                                    server_handler.DiscordMsg_Callback_Failed_0(msgId, parentMsgId, '')
                            gExecutor.submit(next_task, prompt, conversation_id, msgId, parentMsgId)
                            response_data=dict(code=0,msg='')
                            
                        else:
                            id, response_data = server_handler.SendDiscordMessageMJGenerate(conversation_id, msgId, parentMsgId, prompt)
                       
                    

                elif type == 'upscale':
                    discordMsgId = data.get('discordMsgId', None)
                    msgHash = data.get('msgHash', None)
                    index = data.get('index', None)
                    if index is  None or discordMsgId is None or msgHash is None  or msgId == None:
                        code = -1
                        msg = f'upscale:参数无效 index:{index} discordMsgId:{discordMsgId} msgHash:{msgHash} msgId:{msgId}'
                        response_data=dict(code=code, msg=msg)
                    else:
                        response_data = server_handler.SendDiscordMessageMJUpscale(msgId, parentMsgId, discordMsgId, msgHash, index)
                        

                elif type == 'variation':
                    discordMsgId = data.get('discordMsgId', None)
                    msgHash = data.get('msgHash', None)
                    index = data.get('index', None)
                    if index is  None or discordMsgId is None or msgHash is None  or msgId == None:
                        code = -1
                        msg = f'variation:参数无效 index:{index} discordMsgId:{discordMsgId} msgHash:{msgHash} msgId:{msgId}'
                        response_data=dict(code=code, msg=msg)
                    else:
                        response_data = server_handler.SendDiscordMessageMJVariate(msgId, parentMsgId, discordMsgId, msgHash, index)

                   
                elif type == 'maxUpscale':
                    discordMsgId = data.get('discordMsgId', None)
                    msgHash = data.get('msgHash', None)
                    if discordMsgId is None or msgHash is None or msgId == None:
                        code = -1
                        msg = f'maxUpscale:参数无效 discordMsgId:{discordMsgId} msgHash:{msgHash}  msgId:{msgId}'
                        response_data=dict(code=code, msg=msg)
                    else:
                        response_data = server_handler.SendDiscordMessageMJMaxUpscale(msgId, parentMsgId, discordMsgId, msgHash)
                        

                
                elif type == 'reset':
                    discordMsgId = data.get('discordMsgId', None)
                    msgHash = data.get('msgHash', None)
                    if discordMsgId is None or msgHash is None or msgId == None:
                        code = -1
                        msg = f'reset:参数无效 discordMsgId:{discordMsgId} msgHash:{msgHash}'
                        response_data=dict(code=code, msg=msg)
                    else:
                        response_data = server_handler.SendDiscordMessageMJReset(msgId, parentMsgId, discordMsgId, msgHash)
                elif type == 'describe':
                    Logger.info('describe.....')
                    imageData = data.get('imageData', [])
                    imageUrl = data.get('imageUrl', None)
                    conversation_id = data.get('conversation_id', '')
                    if (imageUrl is None and (imageData is None or len(imageData) <= 0)) or msgId is None:
                        code = -1
                        msg = f'describe:图片数据不正确 msgId:{msgId}'
                        response_data=dict(code=-1, msg=msg)
                    else:    
                        def next_task(conversation_id, msgId, imageUrl, imageData):
                            server_handler.SendDiscordMessageDescribe(conversation_id, msgId, imageUrl, imageData)
                        server_handler.gExecutor.submit(next_task, conversation_id, msgId, imageUrl, imageData)
                        response_data=dict(code=0, msg='任务已提交、等待处理中')
                else:
                    response_data=dict(code=-3, msg=f'type invalid {type}')
            except Exception as e:
                Logger.error(str(e), exc_info=True)
                response_data=dict(code=-4, msg=f'err broken')
            
            # 处理POST请求
            
            response_body = json.dumps(response_data).encode('utf-8')
            
            self.send_response(200)
            self.send_header('Content-Type', 'application/json')
            self.send_header('Content-Length', len(response_body))
            self.end_headers()
            self.wfile.write(response_body)
            
        elif discord_callback.match(self.path):

            content_length = int(self.headers['Content-Length'])
            try:
                body = self.rfile.read(content_length)
                data = json.loads(body)
                
                type = data.get('type', None)    
                
                
                if type == 'FirstTrigger':
                    content = data.get('content', '')
                    id,err = server_handler.getDiscordLardMapId(content)
                    #discordIteratorTag(id)
                    Logger.info(f'type == FirstTrigger id:{id}')
                    if err is None:
                        response_data = dict(code=0, msg='')
                    else:
                        response_data=dict(code=-5, msg=f'发生错误 {self.path}')
                    
                elif type == 'GenerateEnd':
                    try:
                        discord = data.get('discord', None)
                        if discord is None:
                            content = ''
                        else:
                            content = discord.get('content', '')

                        
                        imageUrl = discord['attachments'][0]['url']
                        id,err = server_handler.getDiscordLardMapId(content)
                        if err is None:
                            
                            msgHash = generateDiscordMsgHash(imageUrl)
                            message_reference = discord.get('message_reference', None)
                            if message_reference is not None:
                                parent_discord_msgId = message_reference.get('message_id', None)
                            else:
                                parent_discord_msgId = None
                            discordMsgId = discord.get('id','')
                            isOk, urls = mj_util.update_discord_user_result_image(imageUrl)
                            if isOk:
                                response_data=server_handler.DiscordMsg_Callback_OnImageEnd(id, urls[0], urls[1], discordMsgId, msgHash, parent_discord_msgId)
                            else:
                                response_data=server_handler.DiscordMsg_Callback_OnImageEnd(id, imageUrl,imageUrl, discordMsgId, msgHash, parent_discord_msgId)
                        else:
                            response_data=dict(code=-5, msg=f'发生错误 {self.path}')

                    except Exception as e:
                        traceback.print_exc()
                        Logger.error(str(e), exc_info=True)
                        Logger.info(f'发生错误 content:{content}')
                        response_data=dict(code=-5, msg=f'发生错误 {self.path}')
                        
                elif type == 'GenerateEditError':
                    try:
                        discord = data.get('discordupdate', None)
                        id,err = server_handler.getDiscordLardMapId(content)
                        if err is None and discord is not None:
                            message_reference = discord.get('message_reference', None)
                            if message_reference is not None:
                                parent_discord_msgId = message_reference.get('message_id', None)
                            else:
                                parent_discord_msgId = None
                            referenced_message = discord.get('referenced_message',None)
                            if referenced_message is None:
                                discordMsgId = None
                            else:
                                discordMsgId = referenced_message.get('id', None)
                            server_handler.DiscordMsg_Callback_Failed(id, discordMsgId, parent_discord_msgId)
                    except Exception as e:
                        traceback.print_exc()
                        response_data=dict(code=-5, msg=f'🤖️ ：发送的指令存在异常，请检查后重试 {self.path}')
                elif type == 'RichText':
                    try:
                        embeds = data.get('embeds', [])
                        embed = embeds[0]
                        if embed:
                            description = embed.get('description', '')
                            imageUrl = embed.get('image',{}).get('url', None)
                            id = None
                            if imageUrl:
                                first_index = imageUrl.rfind('/')
                                if first_index >= 0:
                                    id = imageUrl[first_index+1:]
                                
                            if id:
                                server_handler.DiscordMsg_Callback_OnDescribeSuccess(id, description)
                    except:
                        traceback.print_exc()

                else:
                    Logger.info(f'type == {type}')
                    response_data=dict(code=-5, msg=f'invalid path {self.path}')
            except Exception as e:
                Logger.error(str(e), exc_info=True)
                response_data=dict(code=-4, msg=f'discord callback: type invalid {type}')
            
            # 处理POST请求
            response_body = json.dumps(response_data).encode('utf-8')
            self.send_response(200)
            self.send_header('Content-Type', 'application/json')
            self.send_header('Content-Length', len(response_body))
            self.end_headers()
            self.wfile.write(response_body)

        else:
            response_data=dict(code=-5, msg=f'path invalid {type}')
             # 处理POST请求
            response_body = json.dumps(response_data).encode('utf-8')

            self.send_response(200)
            self.send_header('Content-Type', 'application/json')
            self.send_header('Content-Length', len(response_body))
            self.end_headers()
            self.wfile.write(response_body)

       

# 定义多线程Web服务器
class ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
    pass






http_server = None
def start_servers(port):
    try:
        global http_server
        httpd = ThreadedHTTPServer(('0.0.0.0', port), MyRequestHandler)
        # 创建 HTTP Server 任务
        httpd.serve_forever()
        
        
    except Exception as e:
        print(f'Server error: {e}')
    finally:
        if http_server:
            http_server.shutdown()
            http_server.server_close()
        #executor.wait(False)




def start(port=16088):
# 启动服务器
    print('start_servers..')
    server_handler.init_timer_work()
    start_servers(port)
    